#include <vector>
#include <memory>
#include "branch_network.h"
#include "trunk_network.h"
#include "loss_function.h"
#include <iostream>

class DeepONet {
public:
    BranchNetwork branchNet;
    TrunkNetwork trunkNet;
    std::shared_ptr<ILossFunction> lossFunction;

    /**
     * 构造函数初始化DeepONet
     *
     * @param branchNet 分支网络
     * @param trunkNet 主干网络
     * @param lossFunc 损失函数
     */
    DeepONet(BranchNetwork branchNet, TrunkNetwork trunkNet, std::shared_ptr<ILossFunction> lossFunc)
        : branchNet(branchNet), trunkNet(trunkNet), lossFunction(lossFunc) {}

    /**
     * 训练DeepONet
     *
     * @param branchInputs 分支网络的输入
     * @param trunkInputs 主干网络的输入
     * @param expectedOutputs 期望的输出
     */
    void train(const std::vector<std::vector<float>>& branchInputs, const std::vector<std::vector<float>>& trunkInputs, const std::vector<std::vector<float>>& expectedOutputs) {
        // 前向传播
        auto branchOutputs = branchNet.forward(branchInputs);
        auto trunkOutputs = trunkNet.forward(trunkInputs);

        // 输出合并和损失计算
        auto combinedOutputs = combineOutputs(branchOutputs, trunkOutputs);
        auto loss = lossFunction->loss(combinedOutputs, expectedOutputs);

        // 输出损失值
        std::cout << "Current Loss: " << loss << std::endl;
        
        // 反向传播
        auto lossGradients = lossFunction->gradient(combinedOutputs, expectedOutputs);
        branchNet.backward(lossGradients);
        trunkNet.backward(lossGradients);
    }

    /**
     * 预测函数
     *
     * @param branchInputs 分支网络的输入
     * @param trunkInputs 主干网络的输入
     * @return 预测的输出
     */
    std::vector<std::vector<float>> predict(const std::vector<std::vector<float>>& branchInputs, const std::vector<std::vector<float>>& trunkInputs) {
        auto branchOutputs = branchNet.forward(branchInputs);
        auto trunkOutputs = trunkNet.forward(trunkInputs);
        return combineOutputs(branchOutputs, trunkOutputs);
    }

private:
    /**
     * 合并分支网络和主干网络的输出
     *
     * @param branchOutputs 分支网络的输出
     * @param trunkOutputs 主干网络的输出
     * @return 合并后的输出
     */
    // std::vector<std::vector<float>> combineOutputs(const std::vector<std::vector<float>>& branchOutputs, const std::vector<std::vector<float>>& trunkOutputs) {
    //     std::vector<std::vector<float>> combinedOutputs(branchOutputs.size(), std::vector<float>(branchOutputs[0].size()));

    //     for (size_t i = 0; i < branchOutputs.size(); ++i) {
    //         for (size_t j = 0; j < branchOutputs[i].size(); ++j) {
    //             combinedOutputs[i][j] = branchOutputs[i][j] * trunkOutputs[i][0];  // Assume trunk output is a scalar factor
    //         }
    //     }
    //     return combinedOutputs;
    // }

    // 修改合并函数为加法示例
    std::vector<std::vector<float>> combineOutputs(const std::vector<std::vector<float>>& branchOutputs, const std::vector<std::vector<float>>& trunkOutputs) {
        std::vector<std::vector<float>> combinedOutputs(branchOutputs.size(), std::vector<float>(branchOutputs[0].size()));

        for (size_t i = 0; i < branchOutputs.size(); ++i) {
            for (size_t j = 0; j < branchOutputs[i].size(); ++j) {
                combinedOutputs[i][j] = branchOutputs[i][j] + trunkOutputs[i][0];  // 使用加法而非乘法
            }
        }
        return combinedOutputs;
    }


};
